# -*- coding: utf-8 -*-
# @Time : 2020/11/25 16:23
# @Author : Jiangzhesheng
# @File : run.py
# @Software: PyCharm
import sys
import argparse
import os
import logging
import myclass
import algo
import subprocess
class Anchor_unit:
    def __init__(self,line='') -> None:
        if line:
            line=line.split()
            self.query_id=line[0]
            self.ref_id=line[3]
            self.ref_start=int(line[4])
            self.ref_end=int(line[5])
            self.orient=line[6]
            self.query_start = int(line[1])
            self.query_end = int(line[2])
            self.type=line[7]

    def tostring(self, ref_first_flag)->str:
        if ref_first_flag:
            info = [self.ref_id,
                    str(self.ref_start),
                    str(self.ref_end),
                    self.query_id,
                    str(self.query_start),
                    str(self.query_end),
                    self.orient,
                    ]
            return '\t'.join(info)
        else:
            info = [self.query_id,
                    str(self.query_start),
                    str(self.query_end),
                    self.ref_id,
                    str(self.ref_start),
                    str(self.ref_end),
                    self.orient,
                    ]
            return '\t'.join(info)

    def __str__(self) -> str:
        return self.tostring(ref_first_flag=True)

def anchor_binning(scriptpath:str,queryfastapath: str, reffastapath:str,outputdir:str):
    os.chdir(os.path.dirname(scriptpath))
    outfilepath=os.path.join(outputdir,'anchor')
    cmd=['bash',scriptpath,reffastapath,queryfastapath,outfilepath]
    a=subprocess.Popen(cmd)
    a.wait()
    return outfilepath

def script_get_edit_distance(script:[(int,str)],start:int,end:int)->int:
    cnt = 0
    for i in script:
        if i[0] >= start and i[0] < end:
            cnt+=1
    return cnt

def script_change_len(script:[(int,str)],start:int,end:int)->int:
    cnt=0
    for i in script:
        if i[0]>=start and i[0]<end:
            if 'D' in i[1]:
                cnt-=1
            elif 'I' in i[1]:
                cnt+=1
    return cnt

def gap_binning(anchor:Anchor_unit, script, logger,pace:int=1000)->[]:
    ans=[]
    if anchor.orient=='+':
        tmp_ref_start,tmp_ref_end=anchor.ref_start,0
        tmp_query_start=anchor.query_start
        i=0
        while True:

            script_start=i*pace
            script_end=(i+1)*pace

            change_length=script_change_len(script=script,start=script_start,end=script_end)
            edit_distance=script_get_edit_distance(script=script,start=script_start,end=script_end)

            tmp_ref_end=tmp_ref_start+pace

            if tmp_ref_end>=anchor.ref_end:
                new_anchor = Anchor_unit()
                new_anchor.ref_id = anchor.ref_id
                new_anchor.ref_start = tmp_ref_start
                new_anchor.ref_end = anchor.ref_end
                new_anchor.query_id = anchor.query_id
                new_anchor.query_start = tmp_query_start
                new_anchor.query_end = anchor.query_end
                new_anchor.orient = anchor.orient

                bin_length = new_anchor.ref_end-new_anchor.ref_start
                idy = (bin_length - edit_distance) / bin_length
                ans.append((new_anchor,idy,edit_distance))
                break

            new_anchor=Anchor_unit()
            new_anchor.ref_id=anchor.ref_id
            new_anchor.ref_start=tmp_ref_start
            new_anchor.ref_end=tmp_ref_end
            new_anchor.query_id=anchor.query_id
            new_anchor.query_start=tmp_query_start
            new_anchor.query_end=new_anchor.query_start+pace+change_length
            new_anchor.orient = anchor.orient

            bin_length=new_anchor.ref_end-new_anchor.ref_start
            idy=(bin_length-edit_distance)/bin_length

            ans.append((new_anchor,idy,edit_distance))

            tmp_ref_start=tmp_ref_start+pace
            tmp_query_start=new_anchor.query_end
            i+=1

    elif anchor.orient=='-':
        tmp_ref_start,tmp_ref_end=anchor.ref_start,0
        tmp_query_start=anchor.query_start
        i=0
        while True:

            script_start=i*pace
            script_end=(i+1)*pace

            change_length=script_change_len(script=script,start=script_start,end=script_end)
            edit_distance=script_get_edit_distance(script=script,start=script_start,end=script_end)

            tmp_ref_end=tmp_ref_start+pace

            if tmp_ref_end>=anchor.ref_end:
                new_anchor = Anchor_unit()
                new_anchor.ref_id = anchor.ref_id
                new_anchor.ref_start = tmp_ref_start
                new_anchor.ref_end = anchor.ref_end
                new_anchor.query_id = anchor.query_id
                new_anchor.query_start = tmp_query_start
                new_anchor.query_end = anchor.query_end
                new_anchor.orient=anchor.orient

                bin_length = new_anchor.ref_end-new_anchor.ref_start
                idy = (bin_length - edit_distance) / bin_length
                ans.append((new_anchor,idy,edit_distance))
                break

            new_anchor=Anchor_unit()
            new_anchor.ref_id=anchor.ref_id
            new_anchor.ref_start=tmp_ref_start
            new_anchor.ref_end=tmp_ref_end
            new_anchor.query_id=anchor.query_id
            new_anchor.query_start=tmp_query_start
            new_anchor.query_end=new_anchor.query_start-pace+change_length#这里和+不一样
            new_anchor.orient = anchor.orient

            bin_length=new_anchor.ref_end-new_anchor.ref_start
            idy=(bin_length-edit_distance)/bin_length

            ans.append((new_anchor,idy,edit_distance))

            tmp_ref_start=tmp_ref_start+pace
            tmp_query_start=new_anchor.query_end
            i+=1
    else:
        logger.error(str(anchor)+' has no orient')

    return ans
def run(queryfastapath: str, reffastapath:str,anchorpath:str,outputdir: str, logger):
    with open(anchorpath,mode='r') as anchorfile,\
            open(reffastapath,mode='r') as reffasta, \
            open(os.path.join(outputdir, os.path.basename(anchorpath) + '_empty.bed'), mode='w') as empty, \
            open(os.path.join(outputdir, os.path.basename(anchorpath)+'-result'), mode='w') as out:

        origin_query=myclass.Fastafile(queryfastapath)
        query_delblank_filepath = myclass.Fastafile.delete_linebreak(origin_query, outpath=outputdir)

        query_d={}
        for i in myclass.Fastafile(query_delblank_filepath):
            query_d[i.id]=i
        i=0
        ref_fasta_unit=myclass.Fasta_unit(reffasta.read())
        for line in anchorfile:
            i+=1
            anchor_unit=Anchor_unit(line)
            query_fasta_unit=query_d[anchor_unit.query_id]
            if anchor_unit.orient=='+':
                query_string=query_fasta_unit.getslice(start=anchor_unit.query_start,end=anchor_unit.query_end)
            elif anchor_unit.orient=='-':
                query_string = query_fasta_unit.getslice(start=anchor_unit.query_end, end=anchor_unit.query_start)
                query_string=myclass.DNA_reverse(query_string)

            ref_string=ref_fasta_unit.getslice(start=anchor_unit.ref_start,end=anchor_unit.ref_end)
            if query_string=='':
                logger.error('query is empty '+str(anchor_unit))
                bed_1 = anchor_unit.query_id + '\t' + \
                      str(anchor_unit.query_start) + '\t' + \
                      str(anchor_unit.query_end) + '\t' + \
                      str(anchor_unit.orient) + '\n'
                bed_2 = anchor_unit.query_id + '\t' + \
                        str(anchor_unit.query_end) + '\t' + \
                        str(anchor_unit.query_start) + '\t' + \
                        str(anchor_unit.orient) + '\n'
                bed=bed_1 if anchor_unit.orient=='+' else bed_2
                empty.write(bed)
                continue

            query_string=query_string.upper()
            ref_string=ref_string.upper()
            if anchor_unit.type=='gap':
                if query_string.count('N')/len(query_string) > 0.3:
                    anchor_unit.orient='N'
                    idy='1'
                    out.write(str(anchor_unit)+'\t'+
                              idy+'\t'+
                              anchor_unit.type+'\t'+
                              str(query_string.count('N')/len(query_string))+'\n')
                    logger.warning('gap has to many N , N fraction %.2f '
                                   %(query_string.count('N')/len(query_string)) +str(anchor_unit))
                    continue
                if max(len(query_string),len(ref_string))/min(len(query_string),len(ref_string)) >1.3:
                    idy = '0'
                    out.write(str(anchor_unit)+'\t'+
                              idy+'\t'+
                              anchor_unit.type+'\t'+
                              str(max(len(query_string),len(ref_string))/min(len(query_string),len(ref_string)))+'\n')
                    logger.warning('gap is unbalance query is %dbp and ref is %dbp '
                                   %(len(query_string),len(ref_string))+str(anchor_unit))
                    continue
                logger.info('gap may process N fraction %.2f , query is %dbp and ref is %dbp '%(query_string.count('N')/len(query_string),len(query_string),len(ref_string))+str(anchor_unit))
            try:
                ond_v=algo.my_ond_algo_version1(A=query_string,B=ref_string)
                edit_dist=algo.get_edit_distance(ond_v)
            except IndexError:
                logger.error(str(anchor_unit)+' failed')
                out.write(str(anchor_unit)+'\n')
                continue
            if anchor_unit.type=='anchor':
                idy=(len(ref_string)-edit_dist)/len(ref_string)
                line=str(anchor_unit)+'\t'+str(idy)+'\t'+str(edit_dist)+'\t'+anchor_unit.type+'\t'+str(query_string.count('N'))+'\n'
                # line = str(anchor_unit) + '\t' + str(idy) + '\t' + str(
                #     edit_dist) + '\t' + anchor_unit.type + '\n'
                out.write(line)
            elif anchor_unit.type=='gap':
                pass
                gap_bin=gap_binning(anchor=anchor_unit, script=algo.get_edit_script(ond_v),logger=logger,pace=2000)
                for gap_anchor in gap_bin:
                    line = str(gap_anchor[0]) + '\t' + str(gap_anchor[1]) + '\t' + str(gap_anchor[2]) +'\t'+anchor_unit.type+ '\n'
                    out.write(line)
            logger.info(str(anchor_unit)+' processed')

def main(argv):
    parser = argparse.ArgumentParser(description='')
    parser.add_argument('-ref', '--reffasta',required=True)
    parser.add_argument('-query', '--queryfasta',required=True)
    parser.add_argument('-anchor', '--anchorpath',required=True)
    parser.add_argument('-o', '--outputdir', default='test_data')
    args = parser.parse_args(argv[1:])

    reffasta=os.path.abspath(args.reffasta)
    queryfasta=os.path.abspath(args.queryfasta)
    anchorpath=os.path.abspath(args.anchorpath)
    outputdir = args.outputdir

    try:
        os.mkdir(outputdir)
    except FileExistsError as e:
        print(outputdir,' is exist, files in it may be overwritten')

    outputdir=os.path.abspath(outputdir)

    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    sh = logging.StreamHandler()
    shfmt = logging.Formatter('%(asctime)s-%(message)s')
    sh.setFormatter(fmt=shfmt)
    sh.setLevel(logging.CRITICAL)
    logger.addHandler(sh)

    fh = logging.FileHandler(filename=os.path.join(outputdir,os.path.basename(anchorpath)+'.log'),mode='w')
    fhfmt = logging.Formatter('%(asctime)s - %(pathname)s[line:%(lineno)d] - %(levelname)s: %(message)s')
    fh.setFormatter(fmt=fhfmt)
    fh.setLevel(logging.INFO)
    logger.addHandler(fh)

    logger.info('cmd:'+str(argv)+'\n')

    run(queryfastapath=os.path.abspath(queryfasta),
        reffastapath=os.path.abspath(reffasta),
        anchorpath=os.path.abspath(anchorpath),
        outputdir=os.path.abspath(outputdir),
        logger=logger)


if __name__ == '__main__':
    main(sys.argv)
